'''
createDate: 2024-03-26 13:02:31
Author: liangc
LastEditTime: 2024-03-26 14:34:43
LastEditors: zclee
FilePath: /quant/ProgramEntrance.py
Description: 
'''

from ppq.api import *
from ppq import *
import torch
import torchvision 
from torch.utils.data import DataLoader

DEVICE = 'cuda'
# 量化规则
PLATFORM = TargetPlatform.PPL_CUDA_INT8

# 数据整理函数
def collate_fn(batch: torch.Tensor) -> torch.Tensor:
    return batch.to(DEVICE)

# 一般是从训练数据进行抽样 1024个样本 比较合适 
# calibration_dataset
dataset = [torch.rand(size=(3, 244, 244)) for _ in range(1024)]

model = torchvision.models.mobilenet.mobilenet_v2() 
model = model.to(DEVICE)

quant_setting = QuantizationSettingFactory.pplcuda_setting()
# quant_setting.advanced_opti

calibration_dataloader = DataLoader(dataset=dataset, batch_size=32)

# 量化模型
quantized = quantize_torch_model(
    model=model,
    calib_dataloader=calibration_dataloader,
    calib_steps=32,     # 做多少次calibration 回覆盖dataloader的bath数量 必须是8——512的数字
    input_shape=[1, 3, 224, 224],   # 拿这个输入来追踪样本信息的
    # inputs=calibration_dataloader[0]  # 如果输入是多尺寸的，直接拿数据集的真是数据传入即可
    setting=quant_setting,
    collate_fn=collate_fn,
    platform=PLATFORM,
    onnx_export_file='model.onnx',  # 量化完模型的名称 因为torch -> onnx 是没有跑过量化的
    device=DEVICE, 
    verbose=1
)

assert isinstance(quantized, BaseGraph)

# 导出量化图
export_ppq_graph(graph=quantized, 
                 platform=PLATFORM, # 导出到什么平台上去
                 graph_save_to='Output/quantized.onnx', # 模型导出到哪里
                 config_save_to='Output/quantized.json')    # 量化参数导出到哪里
# 量化结束

# 分析量化误差
reports = graphwise_error_analyse(
    graph=quantized,
    running_device=DEVICE,
    collate_fn=collate_fn,
    dataloader=calibration_dataloader
)

